import time

import pandas as pd
import torch
from torch import nn, optim


class MaxStateSuper(torch.nn.Module):
    def __init__(self, dim_size, heads):
        super(MaxStateSuper, self).__init__()
        self.heads = heads
        assert dim_size % heads == 0, "Dimension size must be divisible by head size."
        self.combined = nn.Linear(dim_size, 4 * dim_size, bias=False)
        self.alpha1 = torch.nn.Parameter(torch.tensor(0.5))
        self.alpha2 = torch.nn.Parameter(torch.tensor(0.5))
        self.alpha3 = torch.nn.Parameter(torch.tensor(0.5))
        self.alpha4 = torch.nn.Parameter(torch.tensor(0.5))

    def forward(self, x, state=None):
        b, s, d = x.shape
        combined = self.combined(x).view(b, s, 4, self.heads, -1)
        out, out1, out2, out3 = combined.unbind(2)
        out = out.permute(0, 3, 1, 2)
        out1 = out1.permute(0, 3, 1, 2)
        out2 = out2.permute(0, 3, 1, 2)
        out3 = out3.permute(0, 3, 1, 2)
        if  state is None:

            out4, _ = torch.cummax(out2, dim=2)
            state = out4[:, :, -1:]
        else:
            out4,_ = torch.cummax(torch.cat([state,out2], dim=2), dim=2)
            state = out4[:, :, -1:]
            out4 = out4[:, :, -1:]
        out = self.gen_model(out, out1, out2, out3, out4)
        out = out.transpose(1, 2).contiguous().view(b, s, d)

        return out, state

    def gen_model(self, a, b, c, d, e):
        term1 = a * b
        term2 = self.alpha1 * b + self.alpha2 * d
        term3 = a * (self.alpha3 * e + d)
        term4 = b * (c + e)
        return term1 + term2 + term3 + term4 + c * e

    # def gen_model(self, a, b, c, d, e):
    #     x = self.alpha1 * b + self.alpha2*d +a
    #     x = a*b + x
    #     x = self.alpha3 *a*e+x
    #     x = b*c +x
    #     x= b*e  +x
    #     x= c*e +x
    #     x= a*d+x
    #
    #
    #     # ab = a * b
    #     # ad = a * d
    #     # ae = a * e
    #     # bc = b * c
    #     # be = b * e
    #     # ce = c * e
    #
    #     # # # 初始计算
    #
    #     # x = self.layer_norm(ab + 2 * b )+ 2 * ae + 2 * d + bc + be + a + self.layer_norm(ce + ad)
    #
    #     return x
    #


class FeedForward(torch.nn.Module):
    def __init__(self, hidden_size):
        super(FeedForward, self).__init__()
        self.ffn1 = torch.nn.Linear(hidden_size, hidden_size)
        self.ffn2 = torch.nn.Linear(hidden_size, hidden_size)
        self.gate = torch.nn.Linear(hidden_size, hidden_size)

        self.relu = torch.nn.ReLU()
        # self.gr = torch.nn.Dropout(0.02)

    def forward(self, x):
        x1 = self.ffn1(x)
        x2 = self.relu(self.gate(x))
        xx = x1 * x2
        x = self.ffn2(xx)
        return x


class DecoderLayer(torch.nn.Module):
    def __init__(self, hidden_size, num_heads):
        super(DecoderLayer, self).__init__()
        self.self_attention = MaxStateSuper(hidden_size, num_heads)
        # self.self_attention = MaxState(hidden_size, num_heads)
        self.ffn = FeedForward(hidden_size)
        self.layer_norm = torch.nn.LayerNorm(hidden_size)

        self.alpha = torch.nn.Parameter(torch.tensor(0.5))

    def forward(self, x, state=None, ):
        x1, state = self.self_attention(x, state)
        x = self.layer_norm(self.alpha * self.ffn(x1) + (1 - self.alpha) * x)

        return x, state


class SamOut(torch.nn.Module):
    def __init__(self, voc_size, hidden_size, num_heads, num_layers):
        super(SamOut, self).__init__()
        self.em = torch.nn.Embedding(voc_size, hidden_size, padding_idx=0)

        self.decoder_layers = torch.nn.ModuleList([DecoderLayer(hidden_size, num_heads) for _ in range(num_layers)])
        self.head = nn.Linear(hidden_size, voc_size, bias=False)
        # self.alpha = [torch.nn.Parameter(torch.tensor(0.5)) for i in range(num_layers)]
        # self.layer_norm = torch.nn.LayerNorm(hidden_size)

    def forward(self, x, state=None):
        x = self.em(x)

        if state is None:
            state = [None] * len(self.decoder_layers)
        i = 0

        for ii, decoder_layer in enumerate(self.decoder_layers):
            x1, state[i] = decoder_layer(x, state[i])
            x = x1 + x
            i += 1

        x = self.head(x)

        return x, state


if __name__ == '__main__':

    # 这里假设 DecoderLayer 已经定义好了，具体实现可以参考之前提供的代码或根据需要自定义

    # 定义超参数
    voc_size = 12506
    num_layers = 8
    hidden_size = 2 ** 6 * num_layers
    num_heads = num_layers
    learning_rate = 0.001
    batch_size = 32
    num_epochs = 1000

    # 初始化模型
    model = SamOut(voc_size=voc_size, hidden_size=hidden_size, num_heads=num_heads, num_layers=num_layers)
    params = 0
    # [i.shape[0]  and len(i.shape) == 1 elif i.shape[1] * i.shape[0]
    for i in model.parameters():
        if i.shape != torch.Size([]):
            params += i.numel()
    print(params)
    # 定义损失函数和优化器
    criterion = nn.CrossEntropyLoss(ignore_index=3)  # 忽略填充标记的损失计算
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)

    # 模拟一些训练数据（实际应用中应该使用真实的数据集）

    # 训练循环
    start_time = time.time()
    for epoch in range(num_epochs):
        data = torch.randint(low=0, high=voc_size, size=(batch_size, 50))  # 输入序列长度为50
        input_tensor = data[:, :-1]
        target_tensor = data[:, 1:]

        # 前向传播
        output, _ = model(input_tensor)

        # 将输出reshape以适应 CrossEntropyLoss 的输入要求
        output = output.reshape(-1, voc_size)
        target_tensor = target_tensor.reshape(-1)

        # 计算损失
        loss = criterion(output, target_tensor)
        # output_mean = (torch.nn.functional.softmax(output, -1)-1).mean()**2
        # c = loss.item() / 50
        # loss = loss - output_mean
        # loss = los
        optimizer.zero_grad()  # 清除梯度

        # 反向传播和优化
        loss.backward()
        optimizer.step()

        print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}--')

    print("Training complete.{}".format(time.time() - start_time))
